{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 分析需求：\n",
    "需要在没有标注的情况下，对考古中的金相图片进行聚类\n",
    "# 实现规划\n",
    "结合天然正例、无监督数据增强，实现图片向量编码，并使用聚类方法对图片进行聚类。\n",
    "# 实现步骤：\n",
    "1. 获取数据集：使用数据集类加载考古中的金相图片，并将多种分辨率版本的相同图片组织配对，构成天然正例。\n",
    "2. 对数据进行预处理，并使用常规数据增强来生成正例。\n",
    "3. 使用simclr模型使用对比学习同时借助监督和无监督方式学习图片向量编码。\n",
    "4. 使用合适的聚类方法对图片进行聚类。\n",
    "\n",
    "# 可能的改进\n",
    "* 查找更先进的对比学习方法。"
   ],
   "id": "de6b6ec2582cce20"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 给大模型的prompt\n",
    "写一个项目，给出实现代码。使用simclr对比学习来通过监督和无监督结合的方式获得图片向量编码，每一个图片的正例包含不同分辨率下的相同物体图片，以及其它通过图片数据增强方式获得的图片，负例采用随机负采样方式获得。之后使用合适的聚类方法对图片进行聚类。"
   ],
   "id": "c9375913319a4f77"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "# 多分辨率生成（固定4种分辨率）\n",
    "RESOLUTIONS = [(256, 256), (128, 128), (64, 64), (32, 32)]\n",
    "\n",
    "# 无监督数据增强（用于生成正例）\n",
    "class SimCLRDataTransform:\n",
    "    def __init__(self, resolution):\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.RandomResizedCrop(resolution),\n",
    "            transforms.RandomHorizontalFlip(p=0.5),\n",
    "            transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),\n",
    "            transforms.RandomGrayscale(p=0.2),\n",
    "            transforms.GaussianBlur(kernel_size=int(0.1 * resolution[0]) + 1, sigma=(0.1, 2.0)),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "        ])\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return [self.transform(x), self.transform(x)]  # 生成2个增强视图\n",
    "\n",
    "# 数据集类（融合多分辨率和增强样本）\n",
    "class MultiResDataset(Dataset):\n",
    "    def __init__(self, root_dir, split='train'):\n",
    "        self.root_dir = root_dir\n",
    "        self.split = split\n",
    "        self.classes = os.listdir(root_dir)\n",
    "        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}\n",
    "        self.samples = self._load_samples()\n",
    "        self.base_transform = transforms.ToTensor()\n",
    "\n",
    "    def _load_samples(self):\n",
    "        samples = []\n",
    "        for cls in self.classes:\n",
    "            cls_dir = os.path.join(self.root_dir, cls)\n",
    "            for img_name in os.listdir(cls_dir):\n",
    "                img_path = os.path.join(cls_dir, img_name)\n",
    "                # 存储原始图片路径+类别ID\n",
    "                samples.append((img_path, self.class_to_idx[cls]))\n",
    "        return samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path, label = self.samples[idx]\n",
    "        img = Image.open(img_path).convert('RGB')\n",
    "        \n",
    "        # 生成多分辨率版本\n",
    "        multi_res_imgs = []\n",
    "        for res in RESOLUTIONS:\n",
    "            resized = img.resize(res, Image.BILINEAR)\n",
    "            multi_res_imgs.append(resized)\n",
    "        \n",
    "        # 对每个分辨率生成增强样本（训练时）\n",
    "        if self.split == 'train':\n",
    "            augmented = []\n",
    "            for res_img in multi_res_imgs:\n",
    "                transform = SimCLRDataTransform(res_img.size)\n",
    "                aug1, aug2 = transform(res_img)\n",
    "                augmented.extend([aug1, aug2])\n",
    "            return {'images': augmented, 'label': label}\n",
    "        else:\n",
    "            # 测试时只返回原始分辨率特征\n",
    "            img_tensor = self.base_transform(img.resize(RESOLUTIONS[0]))\n",
    "            return {'image': img_tensor, 'label': label}\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "# 数据加载器\n",
    "def get_dataloader(root_dir, batch_size=32, split='train'):\n",
    "    dataset = MultiResDataset(root_dir, split)\n",
    "    return DataLoader(\n",
    "        dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=(split == 'train'),\n",
    "        num_workers=4,\n",
    "        pin_memory=True\n",
    "    )"
   ],
   "id": "542b4d96649e0eb8"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 分析存在问题\n",
    "1. 数据增强方式使用工具没有进行注释，需要进行注释。\n",
    "2. 对于多分辨率版本，虽然常规操作是由自己生成，但是这里我们由相应数据，所以应该配对和使用自己的数据进行配对。\n",
    "3. 对于图片文件的读取需要根据实际情况进行修改。这里采用todo形式标记"
   ],
   "id": "c7400c33ef8071c0"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision.models import resnet50\n",
    "\n",
    "# 投影头（用于对比学习）\n",
    "class ProjectionHead(nn.Module):\n",
    "    def __init__(self, in_dim=2048, hidden_dim=2048, out_dim=128):\n",
    "        super().__init__()\n",
    "        self.layers = nn.Sequential(\n",
    "            nn.Linear(in_dim, hidden_dim),\n",
    "            nn.BatchNorm1d(hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, out_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layers(x)\n",
    "\n",
    "# SimCLR模型（融合监督信号）\n",
    "class SimCLR(nn.Module):\n",
    "    def __init__(self, backbone='resnet50', projection_dim=128):\n",
    "        super().__init__()\n",
    "        # 骨干网络（ResNet50）\n",
    "        self.backbone = resnet50(pretrained=False)\n",
    "        self.backbone.fc = nn.Identity()  # 移除原始全连接层\n",
    "        self.feature_dim = 2048  # ResNet50输出维度\n",
    "        \n",
    "        # 投影头\n",
    "        self.projection_head = ProjectionHead(\n",
    "            in_dim=self.feature_dim,\n",
    "            out_dim=projection_dim\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x: [batch_size * num_views, C, H, W]\n",
    "        features = self.backbone(x)  # [batch_size * num_views, 2048]\n",
    "        projections = self.projection_head(features)  # [batch_size * num_views, 128]\n",
    "        return features, projections\n",
    "\n",
    "# 监督+无监督融合的InfoNCE损失\n",
    "def simclr_loss(projections, labels, temperature=0.07):\n",
    "    \"\"\"\n",
    "    projections: [N, D] 其中N = batch_size * num_views\n",
    "    labels: [batch_size] 原始样本标签（每个样本有num_views个投影）\n",
    "    \"\"\"\n",
    "    batch_size = labels.shape[0]\n",
    "    num_views = projections.shape[0] // batch_size  # 每个样本的视图数\n",
    "    \n",
    "    # 计算余弦相似度矩阵\n",
    "    sim = F.cosine_similarity(\n",
    "        projections.unsqueeze(1), \n",
    "        projections.unsqueeze(0), \n",
    "        dim=2\n",
    "    )  # [N, N]\n",
    "    \n",
    "    # 构建掩码：同一原始样本的不同视图为正例\n",
    "    mask = torch.eye(batch_size, device=projections.device).repeat(num_views, num_views)\n",
    "    mask = mask - torch.eye(N, device=projections.device)  # 排除自身\n",
    "    \n",
    "    # 监督掩码：同一类别的样本也视为正例\n",
    "    label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()  # [B, B]\n",
    "    label_mask = label_mask.repeat(num_views, num_views)  # [N, N]\n",
    "    mask = mask + label_mask  # 融合两种正例掩码\n",
    "    mask = mask.clamp(0, 1)  # 确保掩码非负\n",
    "    \n",
    "    # 计算InfoNCE损失\n",
    "    sim = sim / temperature\n",
    "    exp_sim = torch.exp(sim)\n",
    "    exp_sim = exp_sim * (1 - torch.eye(N, device=projections.device))  # 排除自身\n",
    "    \n",
    "    log_prob = sim - torch.log(exp_sim.sum(dim=1, keepdim=True))\n",
    "    loss = (-mask * log_prob).sum(dim=1) / mask.sum(dim=1)  # 平均正例损失\n",
    "    return loss.mean()"
   ],
   "id": "2ab7e28a0647c13f"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 分析存在问题\n",
    "\n",
    "1. 目前的backbone模型为resnet50，后续可能根据需要更换成ViT，甚至多模态大模型\n",
    "2. SimCLR模型本身没什么问题，主要需要关注simclr_loss函数，"
   ],
   "id": "6de463996222af"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "import torch\n",
    "import argparse\n",
    "import torch.optim as optim\n",
    "from tqdm import tqdm\n",
    "from data_preprocess import get_dataloader\n",
    "from simclr_model import SimCLR, simclr_loss\n",
    "\n",
    "def train(args):\n",
    "    # 设备配置\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    \n",
    "    # 数据加载\n",
    "    train_loader = get_dataloader(\n",
    "        root_dir=args.data_dir,\n",
    "        batch_size=args.batch_size,\n",
    "        split='train'\n",
    "    )\n",
    "    val_loader = get_dataloader(\n",
    "        root_dir=args.data_dir,\n",
    "        batch_size=args.batch_size,\n",
    "        split='val'\n",
    "    )\n",
    "    \n",
    "    # 模型初始化\n",
    "    model = SimCLR(projection_dim=args.projection_dim).to(device)\n",
    "    optimizer = optim.AdamW(\n",
    "        model.parameters(),\n",
    "        lr=args.lr,\n",
    "        weight_decay=args.weight_decay\n",
    "    )\n",
    "    scheduler = optim.lr_scheduler.CosineAnnealingLR(\n",
    "        optimizer, \n",
    "        T_max=args.epochs\n",
    "    )\n",
    "    \n",
    "    # 训练循环\n",
    "    best_loss = float('inf')\n",
    "    for epoch in range(args.epochs):\n",
    "        model.train()\n",
    "        train_loss = 0.0\n",
    "        \n",
    "        for batch in tqdm(train_loader, desc=f'Epoch {epoch+1}'):\n",
    "            images = batch['images']  # list of tensors: [num_views*num_res, B, C, H, W]\n",
    "            labels = batch['label'].to(device)  # [B]\n",
    "            \n",
    "            # 整理输入格式：[B * num_views, C, H, W]\n",
    "            num_views = len(images) // labels.shape[0]\n",
    "            images = torch.cat(images, dim=0).to(device)  # [B*num_views, C, H, W]\n",
    "            \n",
    "            # 前向传播\n",
    "            _, projections = model(images)  # [B*num_views, projection_dim]\n",
    "            \n",
    "            # 计算损失\n",
    "            loss = simclr_loss(projections, labels, temperature=args.temperature)\n",
    "            \n",
    "            # 反向传播\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            train_loss += loss.item() * images.shape[0]\n",
    "        \n",
    "        # 验证\n",
    "        model.eval()\n",
    "        val_loss = 0.0\n",
    "        with torch.no_grad():\n",
    "            for batch in val_loader:\n",
    "                images = batch['image'].to(device)  # [B, C, H, W]\n",
    "                labels = batch['label'].to(device)\n",
    "                _, projections = model(images)\n",
    "                loss = simclr_loss(projections, labels, temperature=args.temperature)\n",
    "                val_loss += loss.item() * images.shape[0]\n",
    "        \n",
    "        # 日志\n",
    "        train_loss /= len(train_loader.dataset)\n",
    "        val_loss /= len(val_loader.dataset)\n",
    "        print(f'Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}')\n",
    "        \n",
    "        # 保存最佳模型\n",
    "        if val_loss < best_loss:\n",
    "            best_loss = val_loss\n",
    "            torch.save(model.backbone.state_dict(), os.path.join(args.save_dir, 'best_backbone.pth'))\n",
    "        \n",
    "        scheduler.step()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--data_dir', type=str, required=True, help='数据集根目录')\n",
    "    parser.add_argument('--save_dir', type=str, default='checkpoints', help='模型保存目录')\n",
    "    parser.add_argument('--projection_dim', type=int, default=128, help='投影头输出维度')\n",
    "    parser.add_argument('--batch_size', type=int, default=32, help='批次大小')\n",
    "    parser.add_argument('--lr', type=float, default=1e-4, help='学习率')\n",
    "    parser.add_argument('--weight_decay', type=float, default=1e-5, help='权重衰减')\n",
    "    parser.add_argument('--temperature', type=float, default=0.07, help='温度参数')\n",
    "    parser.add_argument('--epochs', type=int, default=100, help='训练轮次')\n",
    "    args = parser.parse_args()\n",
    "    \n",
    "    os.makedirs(args.save_dir, exist_ok=True)\n",
    "    train(args)"
   ],
   "id": "f4d547c8e6b4bcbf"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import argparse\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "from sklearn.metrics import (\n",
    "    adjusted_rand_score,\n",
    "    normalized_mutual_info_score,\n",
    "    completeness_score\n",
    ")\n",
    "from data_preprocess import get_dataloader\n",
    "from torchvision.models import resnet50\n",
    "\n",
    "def extract_features(args):\n",
    "    \"\"\"提取所有图片的特征向量\"\"\"\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    \n",
    "    # 加载模型\n",
    "    model = resnet50(pretrained=False)\n",
    "    model.fc = nn.Identity()  # 移除分类头\n",
    "    model.load_state_dict(torch.load(os.path.join(args.save_dir, 'best_backbone.pth')))\n",
    "    model = model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    # 加载数据\n",
    "    dataloader = get_dataloader(\n",
    "        root_dir=args.data_dir,\n",
    "        batch_size=args.batch_size,\n",
    "        split='val'\n",
    "    )\n",
    "    \n",
    "    # 提取特征\n",
    "    features = []\n",
    "    labels = []\n",
    "    with torch.no_grad():\n",
    "        for batch in dataloader:\n",
    "            images = batch['image'].to(device)\n",
    "            feat = model(images)  # [B, 2048]\n",
    "            features.append(feat.cpu().numpy())\n",
    "            labels.append(batch['label'].numpy())\n",
    "    \n",
    "    return np.vstack(features), np.hstack(labels)\n",
    "\n",
    "def cluster_and_evaluate(features, true_labels):\n",
    "    \"\"\"层次聚类并评估\"\"\"\n",
    "    # L2归一化\n",
    "    features = features / np.linalg.norm(features, axis=1, keepdims=True)\n",
    "    \n",
    "    # 层次聚类（基于余弦距离）\n",
    "    clustering = AgglomerativeClustering(\n",
    "        n_clusters=None,\n",
    "        affinity='cosine',\n",
    "        linkage='average',\n",
    "        distance_threshold=0.3  # 距离阈值控制聚类粒度\n",
    "    )\n",
    "    pred_labels = clustering.fit_predict(features)\n",
    "    \n",
    "    # 评估指标\n",
    "    ari = adjusted_rand_score(true_labels, pred_labels)\n",
    "    nmi = normalized_mutual_info_score(true_labels, pred_labels)\n",
    "    completeness = completeness_score(true_labels, pred_labels)\n",
    "    \n",
    "    print(f'Adjusted Rand Index: {ari:.4f}')\n",
    "    print(f'Normalized Mutual Information: {nmi:.4f}')\n",
    "    print(f'Completeness: {completeness:.4f}')\n",
    "    \n",
    "    return pred_labels\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--data_dir', type=str, required=True, help='数据集根目录')\n",
    "    parser.add_argument('--save_dir', type=str, default='checkpoints', help='模型保存目录')\n",
    "    parser.add_argument('--batch_size', type=int, default=32, help='批次大小')\n",
    "    args = parser.parse_args()\n",
    "    \n",
    "    # 提取特征并聚类\n",
    "    features, true_labels = extract_features(args)\n",
    "    cluster_and_evaluate(features, true_labels)"
   ],
   "id": "4571c85b04ed4d29"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    " cluster.py\n",
    "中的数据实际没有labels，将代码修改为只提取特征并聚类，不做评价​"
   ],
   "id": "f3401870842118d1"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "在IDE中让trae添加注释\n",
    "\n",
    " #__init__\n",
    "给使用的所有transform操作做注释\n"
   ],
   "id": "b63228d2fa3751a9"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# 无监督数据增强（用于生成正例）\n",
    "class SimCLRDataTransform:\n",
    "    def __init__(self, resolution):\n",
    "        self.transform = transforms.Compose([\n",
    "            # 随机裁剪并调整大小至指定分辨率，增加样本多样性\n",
    "            transforms.RandomResizedCrop(resolution),\n",
    "            # 以0.5的概率随机水平翻转，模拟不同视角\n",
    "            transforms.RandomHorizontalFlip(p=0.5),\n",
    "            # 以0.8的概率应用颜色抖动，参数分别为亮度、对比度、饱和度和色调的调整范围\n",
    "            transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),\n",
    "            # 以0.2的概率转换为灰度图，保留图像结构信息\n",
    "            transforms.RandomGrayscale(p=0.2),\n",
    "            # 应用高斯模糊，核大小根据分辨率动态调整，sigma在0.1到2.0之间随机\n",
    "            transforms.GaussianBlur(kernel_size=int(0.1 * resolution[0]) + 1, sigma=(0.1, 2.0)),\n",
    "            # 将PIL图像转换为PyTorch张量，范围[0,1]\n",
    "            transforms.ToTensor(),\n",
    "            # 使用ImageNet预训练模型的均值和标准差进行归一化\n",
    "            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "        ])\n"
   ],
   "id": "518be3ed9f737874"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "增加todo标签处理数据预处理",
   "id": "ae57b137e932ccdc"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### 分析\n",
    "* simclr_loss函数和预期不符合，需要大幅修改\n",
    "* 现在正例、负例信息是靠labels获取的，但是我们的数据中缺少labels信息，需要自己处理\n",
    "* 希望能改成直接输入正例对，但是这里就会涉及多个文件的修改\n",
    "\n",
    "### prompt\n",
    " # simclr_loss\n",
    " # data_preprocess.py\n",
    "现在simclr_loss函数的输入是需要labels参数的，而这个参数是由data_preprocess提供的，但是我们的数据集中并不包含这个信息，希望能在data_preprocess中支持在\n",
    " _load_samples\n",
    "中输入正例，并生成labels\n"
   ],
   "id": "6c3840a5566088fc"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "",
   "id": "904b6fa2355c7ad0"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import argparse\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "from data_preprocess import get_dataloader\n",
    "from torchvision.models import resnet50\n",
    "import torch.nn as nn  # 添加缺失的导入\n",
    "\n",
    "\n",
    "def extract_features(args):\n",
    "    \"\"\"提取所有图片的特征向量\"\"\"\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    \n",
    "    # 加载模型\n",
    "    model = resnet50(pretrained=False)\n",
    "    model.fc = nn.Identity()  # 移除分类头\n",
    "    model.load_state_dict(torch.load(os.path.join(args.save_dir, 'best_backbone.pth')))\n",
    "    model = model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    # 加载数据\n",
    "    dataloader = get_dataloader(\n",
    "        root_dir=args.data_dir,\n",
    "        batch_size=args.batch_size,\n",
    "        split='val'\n",
    "    )\n",
    "    \n",
    "    # 提取特征和图像路径\n",
    "    features = []\n",
    "    image_paths = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for batch in dataloader:\n",
    "            images = batch['image'].to(device)\n",
    "            feat = model(images)  # [B, 2048]\n",
    "            features.append(feat.cpu().numpy())\n",
    "            \n",
    "            # 如果数据集提供了图像路径，保存它们\n",
    "            if hasattr(dataloader.dataset, 'samples'):\n",
    "                start_idx = len(image_paths)\n",
    "                end_idx = start_idx + images.shape[0]\n",
    "                # 获取当前批次的图像路径\n",
    "                batch_paths = [sample[0] for sample in dataloader.dataset.samples[start_idx:end_idx]]\n",
    "                image_paths.extend(batch_paths)\n",
    "    \n",
    "    return np.vstack(features), image_paths\n",
    "\n",
    "\n",
    "def perform_clustering(features, image_paths=None, distance_threshold=0.3):\n",
    "    \"\"\"执行层次聚类并返回聚类结果\"\"\"\n",
    "    # L2归一化\n",
    "    features = features / np.linalg.norm(features, axis=1, keepdims=True)\n",
    "    \n",
    "    # 层次聚类（基于余弦距离）\n",
    "    clustering = AgglomerativeClustering(\n",
    "        n_clusters=None,\n",
    "        affinity='cosine',\n",
    "        linkage='average',\n",
    "        distance_threshold=distance_threshold  # 距离阈值控制聚类粒度\n",
    "    )\n",
    "    pred_labels = clustering.fit_predict(features)\n",
    "    \n",
    "    # 获取聚类数量\n",
    "    n_clusters = len(set(pred_labels))\n",
    "    print(f'聚类完成，共得到 {n_clusters} 个簇')\n",
    "    \n",
    "    # 如果提供了图像路径，输出每个簇包含的图像数量\n",
    "    if image_paths is not None:\n",
    "        cluster_counts = {i: 0 for i in range(n_clusters)}\n",
    "        for label in pred_labels:\n",
    "            cluster_counts[label] += 1\n",
    "        \n",
    "        print('每个簇的图像数量:')\n",
    "        for cluster_id, count in sorted(cluster_counts.items()):\n",
    "            print(f'簇 {cluster_id}: {count} 张图像')\n",
    "    \n",
    "    return pred_labels\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--data_dir', type=str, required=True, help='数据集根目录')\n",
    "    parser.add_argument('--save_dir', type=str, default='checkpoints', help='模型保存目录')\n",
    "    parser.add_argument('--batch_size', type=int, default=32, help='批次大小')\n",
    "    parser.add_argument('--distance_threshold', type=float, default=0.3, help='聚类距离阈值')\n",
    "    args = parser.parse_args()\n",
    "    \n",
    "    # 提取特征并聚类\n",
    "    features, image_paths = extract_features(args)\n",
    "    pred_labels = perform_clustering(features, image_paths, args.distance_threshold)\n",
    "    \n",
    "    # 可选：保存聚类结果\n",
    "    np.save(os.path.join(args.save_dir, 'cluster_labels.npy'), pred_labels)\n",
    "    print(f'聚类标签已保存到 {os.path.join(args.save_dir, 'cluster_labels.npy')}')\n",
    "    \n",
    "    if len(image_paths) > 0:\n",
    "        # 保存图像路径和对应聚类标签的映射\n",
    "        with open(os.path.join(args.save_dir, 'image_cluster_mapping.txt'), 'w') as f:\n",
    "            for img_path, cluster_id in zip(image_paths, pred_labels):\n",
    "                f.write(f'{img_path},{cluster_id}\\n')\n",
    "        print(f'图像-聚类映射已保存到 {os.path.join(args.save_dir, 'image_cluster_mapping.txt')}')\n"
   ],
   "id": "dcb6e0d36885668c"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
